import numpy as np
import pandas as pd
import scipy.stats as stats
import itertools
from scipy import special


class StudentPara(object):
    def __init__(self, n):
        self.n = n
        super(object, self).__init__()

    def mul_ability(self, mean, cov):
        return np.random.multivariate_normal(mean=mean, cov=cov, size=self.n)

    def sun_ability(self, ability, beta, ability_variance, rescale=False):

        sun_ability_mu = np.outer(ability, beta)
        sigma = ability_variance - beta ** 2 * ability_variance
        e = np.array([np.random.normal(0, sigma[i] ** (1 / 2), size=self.n) for i in range(len(sigma))]).T
        x = sun_ability_mu + e
        if rescale == True:
            return (x - np.nanmean(x, axis=0)) / np.nanstd(x) * (ability_variance ** (0.5))
        else:
            return x

    def get_att(self, ability, lambda_k, lambda_0, return_att_possibility=False):
        raw_att = np.outer(ability, lambda_k) + lambda_0
        att_possibility = np.exp(raw_att) / (1 + np.exp(raw_att))
        att = pd.DataFrame(att_possibility).applymap(
            lambda x: np.random.binomial(1, x, 1)).values.astype(
            int)
        if return_att_possibility == True:
            return att, att_possibility
        else:
            return att


class ItemPara(object):
    def __init__(self, item_cnt=None, input_Q=None):
        if input_Q is None:
            l1 = [np.eye(4) for i in range(5)]
            l2 = [[1, 1, 0, 0], [0, 1, 1, 0], [0, 0, 1, 1], [1, 0, 0, 1]]
            l3 = [[1, 1, 0, 0], [0, 1, 1, 0], [0, 0, 1, 1], [1, 0, 0, 1]]
            l4 = [[1, 1, 1, 0], [0, 1, 1, 1]]
            self.Q = np.vstack((np.vstack(l1), l2, l3, l4))
            self.i = item_cnt
        else:
            self.Q = input_Q
            self.i = item_cnt
        super(ItemPara, self).__init__()

    def ac_para(self, a, b, c):
        a = np.random.uniform(a[0], a[1], size=self.Q.shape)
        b = np.random.normal(b[0], b[1] ** 0.5, size=self.i)
        c = np.random.uniform(c[0], c[1], size=self.i)
        return a * self.Q, b, c

    def rt_para(self, xi_para, omega_para):
        xi = np.random.normal(xi_para[0], xi_para[1] ** 0.5, size=self.i)
        omega = np.random.normal(omega_para[0], omega_para[1] ** 0.5, size=self.i)
        return xi, omega

    def get_irt_ra(self, theta, a, b, c, return_resp_possibility=False):
        raw_resp_list = np.zeros((len(theta), self.i))
        for i in range(len(a)):
            phi = np.nansum(theta * a[i], axis=1)
            raw_resp_list[:, i] = c[i] + (1 - c[i]) * (phi - b[i])
        correct_possibility = np.exp(raw_resp_list) / (1 + np.exp(raw_resp_list))
        final_response = pd.DataFrame(correct_possibility).applymap(
            lambda x: np.random.binomial(1, x, 1)).values.astype(
            int)
        if return_resp_possibility == True:
            return final_response, correct_possibility
        return final_response

    def s_g_re_para(self, s, g):
        beta_i = special.logit(1 - s) - special.logit(g)
        delta_i = special.logit(g)
        return beta_i, delta_i

    def get_dina_ra(self, att, beta_i, delta_i, return_possibility=False):
        eta = np.zeros((len(att), len(self.Q)))
        for _ in range(len(self.Q)):
            x = att ** self.Q[_]
            eta[:, _] = np.product(x, axis=1)
        raw_resp = eta * beta_i + delta_i
        resp_possibility = np.exp(raw_resp) / (1 + np.exp(raw_resp))
        resp = pd.DataFrame(resp_possibility).applymap(
            lambda x: np.random.binomial(1, x, 1)).values.astype(
            int)
        if return_possibility == True:
            return resp, resp_possibility
        else:
            return resp

    def get_rt(self, xi, omega, tau):
        rt = np.zeros((len(tau), self.i))
        for i in range(self.i):
            phi = np.nansum(tau * self.Q[i], axis=1)
            epsilon_i = np.random.normal(0, omega[i] ** (-2), size=len(tau))
            rt[:, i] = xi[i] - phi + epsilon_i
        x = np.exp(rt).mean()
        return np.exp(rt)

    def rondom_convert(self, ra, rt, rg_percent, rg_ra, rg_t_mean, rg_t_variance):
        rg_index_mat = np.random.binomial(1, rg_percent, ra.shape)
        rg_ra = np.random.binomial(1, rg_ra, ra.shape) * rg_index_mat
        rg_rt = np.exp(np.random.normal(rg_t_mean, rg_t_variance ** 0.5, ra.shape) * rg_index_mat)
        ra = ra * (1 - rg_index_mat) + rg_ra
        rt = rt * (1 - rg_index_mat) + rg_rt

        b = np.product(rt > 0)
        # for n in range(ra.shape[0]):
        #     for i in range(ra.shape[1]):
        #         if rg_index_mat[n,i] == 1:
        #             ra[n,i] = np.random.binomial(1,rg_ra,1)
        #             rt[n,i] = np.random.normal(rg_t_mean,rg_t_variance)

        return ra, rt, rg_index_mat

    def backward_convert(self, ra, rt, time_limit, rg_ra, rg_t_mean, rg_t_variance):
        rg_index_mat = np.zeros((ra.shape))
        for n in range(rt.shape[0]):
            last_index = 0
            while np.nansum(rt[n]) > time_limit:
                rg_index = -1 - last_index
                last_index -= 1
                rt[n, rg_index] = np.exp(np.random.normal(rg_t_mean, rg_t_variance ** (0.5)))
                ra[n, rg_index] = np.random.binomial(1, rg_ra, 1)
                rg_index_mat[n, rg_index] = 1
        return ra, rt, rg_index_mat
